import pymysql
from pymysql.cursors import DictCursor
import os
import logging

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 数据库配置
DB_CONFIG = {
    'host': os.environ.get('DB_HOST', '116.63.51.45'),
    'port': int(os.environ.get('DB_PORT', 27017)),
    'user': os.environ.get('DB_USER', 'root'),
    'password': os.environ.get('DB_PASSWORD', ''),
    'db': os.environ.get('DB_NAME', 'kube-replayer'),
    'charset': 'utf8mb4',
    'cursorclass': DictCursor
}

def get_db_connection():
    """获取数据库连接"""
    try:
        connection = pymysql.connect(
            host=DB_CONFIG['host'],
            port=DB_CONFIG['port'],
            user=DB_CONFIG['user'],
            password=DB_CONFIG['password'],
            db=DB_CONFIG['db'],
            charset=DB_CONFIG['charset'],
            cursorclass=DB_CONFIG['cursorclass']
        )
        return connection
    except Exception as e:
        logger.error(f"数据库连接失败: {e}")
        raise

def execute_query(query, params=None):
    """执行查询并返回结果"""
    connection = get_db_connection()
    try:
        with connection.cursor() as cursor:
            cursor.execute(query, params)
            result = cursor.fetchall()
            return result
    finally:
        connection.close()

def execute_update(query, params=None):
    """执行更新操作并返回影响的行数"""
    connection = get_db_connection()
    try:
        with connection.cursor() as cursor:
            affected_rows = cursor.execute(query, params)
            connection.commit()
            return affected_rows
    finally:
        connection.close()

def execute_insert(query, params=None):
    """执行插入操作并返回最后插入的ID"""
    connection = get_db_connection()
    try:
        with connection.cursor() as cursor:
            cursor.execute(query, params)
            connection.commit()
            return cursor.lastrowid
    finally:
        connection.close()

def execute_transaction(queries_and_params):
    """执行事务，接受查询和参数的列表"""
    connection = get_db_connection()
    try:
        connection.begin()
        with connection.cursor() as cursor:
            results = []
            for query, params in queries_and_params:
                cursor.execute(query, params)
                results.append(cursor.lastrowid)
            connection.commit()
            return results
    except Exception as e:
        connection.rollback()
        logger.error(f"事务执行失败: {e}")
        raise
    finally:
        connection.close()

def init_db():
    """初始化数据库，执行SQL脚本"""
    try:
        # 读取SQL脚本
        with open('db/init_db.sql', 'r') as f:
            sql_script = f.read()
        
        # 连接到MySQL（不指定数据库）
        connection = pymysql.connect(
            host=DB_CONFIG['host'],
            port=DB_CONFIG['port'],
            user=DB_CONFIG['user'],
            password=DB_CONFIG['password'],
            charset=DB_CONFIG['charset']
        )
        
        try:
            with connection.cursor() as cursor:
                # 执行SQL脚本（按语句分割）
                for statement in sql_script.split(';'):
                    if statement.strip():
                        cursor.execute(statement)
                connection.commit()
                logger.info("数据库初始化成功")
        finally:
            connection.close()
    except Exception as e:
        logger.error(f"数据库初始化失败: {e}")
        raise

if __name__ == "__main__":
    # 当直接运行此脚本时，初始化数据库
    init_db()
